import torch
import torchaudio
import librosa
import asyncio
def pre(audio_file_path):
    # 加载音频文件
    waveform, sample_rate = librosa.load(audio_file_path)
    waveform = torch.tensor(waveform).cuda()
    # 16位量化
    waveform_quantized = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=22050)(waveform)
    # 调整音量为-6dB
    max_amplitude = torch.max(torch.abs(waveform_quantized))
    target_amplitude = 0.5 * max_amplitude / 10 ** (-6 / 20)  # -6dB to amplitude
    waveform_quantized = waveform_quantized * (target_amplitude / max_amplitude)
    # 傅立叶变换
    n_fft = int(sample_rate / 22)  # 设置频率分辨率为22Hz
    spectrogram = torch.stft(waveform_quantized, n_fft=n_fft, hop_length=None, win_length=None,
                             window=torch.hann_window(n_fft).cuda(),return_complex=True)
    #spectrogram = torch.mean(spectrogram,dim=1)
    #spectrogram = torch.abs(spectrogram)
    return spectrogram
async def save(i,s,f):
    pickle.dump(s, f)
    print(i, datetime.now())
if __name__ == '__main__':
    import os
    import pickle
    import pandas
    from datetime import datetime
    b = pandas.read_csv(r"D:\old\Desktop\old\animal\barking-emotion-recognition\data\dataset_2.csv")
    '''
    for index in range(604):
        filestr = r"D:/old/Desktop/old/animal/barking-emotion-recognition/data/audioset_audios/" + b["ytid"][
            index] + "_" + str(int(b["start"][index])) + "_" + str(int(b["stop"][index])) + "_cut.mp3"
        if b["label"][index] == 'Happy':
            label = torch.tensor([1, 0, 0]).float()
        elif b["label"][index] == 'Aggressive':
            label = torch.tensor([0, 1, 0]).float()
        else:
            label = torch.tensor([0, 0, 1]).float()
        if os.path.exists(filestr):
            with open(r"D:/emotiondataset/"+str(index),"wb") as f:
                asyncio.run(save(index,[pre(filestr).cpu(), label],f))
    '''
    list=[[],[],[]]
    for index in range(400):
        filestr = r"D:/old/Desktop/old/animal/barking-emotion-recognition/data/audioset_audios/" + b["ytid"][
            index] + "_" + str(int(b["start"][index])) + "_" + str(int(b["stop"][index])) + "_cut.mp3"
        if os.path.exists(filestr):
            if b["label"][index] == 'Happy':
                list[0].append(index)
            elif b["label"][index] == 'Aggressive':
                list[1].append(index)
            else:
                list[2].append(index)
            with open(r"D:/emotiondataset/" + "trainlist", "wb") as f:
                pickle.dump(list, f)
    print("okay")